import komand
from .schema import DownloadReportInput, DownloadReportOutput
# Custom imports below
import json
import base64
import os
from tenable_io.exceptions import TenableIOApiException
from tenable_io.api.scans import ScansApi, ScanExportRequest
from tenable_io.util import wait_until
from komand_tenable_io.util import util



class DownloadReport(komand.Action):
    def __init__(self):
        super(self.__class__, self).__init__(
                name='download_report',
                description='Download the Base64-encoded report generated by a completed scan',
                input=DownloadReportInput(),
                output=DownloadReportOutput())

    def run(self, params={}):
        try:
            scan = self.connection.client.scan_helper.scans(name=params['name'])[0]
            scan_file = '{}.{}'.format(scan.name(), params['file_format'].lower())
            scan.wait_until_stopped()

            t = params['report_type']
            if t == 'Vulnerabilities by Host':
                chapter = ScanExportRequest.CHAPTER_CUSTOM_VULN_BY_HOST
            elif t == 'Vulnerabilities by Plugin':
                chapter = ScanExportRequest.CHAPTER_CUSTOM_VULN_BY_PLUGIN
            elif t == 'Executive Summary':
                chapter = ScanExportRequest.CHAPTER_EXECUTIVE_SUMMARY

            file_format = 'FORMAT_{}'.format(params['file_format'])

            scan.download(scan_file, chapter=chapter, format=getattr(ScanExportRequest, file_format))

            with open(scan_file, 'rb') as r:
                data = r.read()

            report_bytes = str(base64.b64encode(data), 'utf-8')
            os.remove(scan_file)

            return { "scan_report": { "content": report_bytes, "filename": scan_file } }

        except TenableIOApiException as e:
            self.logger.error("An API error occurred: " + str(e))
        except Exception as e:
            self.logger.error("Scan could not be downloaded. Error: " + str(e))

    def test(self):
        conn = self.connection.client
        return { "scan_report": util.folder_verify(conn, self.logger) }
